#include <set>
#include <localization/localization.h>
#include <localization/vision_model.h>

namespace {
struct LampAssociation {
  // Position of the lamp relative to the global map's frame of reference
  double mapped_x;
  double mapped_y;

  // Observed position of the lamp relative to the robot's frame of reference
  double observed_r;
  double observed_theta;

  // Predicted position of the lamp relative to the robot's frame of reference
  double predicted_r;
  double predicted_theta;
};

std::vector<LampAssociation> AssociateLamps(const player_pose2d *mean,
        const gsl_matrix *covariance,
        const std::vector<player_point_2d> detected_lamps,
        const std::vector<player_point_2d> mapped_lamps);

void PrintDebugInfo(const std::vector<player_point_2d> detected_lamps,
        const std::vector<LampAssociation> associations);
}

VisionModel::VisionModel(const std::vector<player_point_2d> &mapped_lamps,
        SyncedLampCentroids *lamp_centroids) :
        mapped_lamps_(mapped_lamps), lamp_centroids_(lamp_centroids) {
  z_predicted_ = gsl_matrix_alloc(2, 1);
  z_observed_ = gsl_matrix_alloc(2, 1);
  H_ = gsl_matrix_alloc(2, 3);
  S_ = gsl_matrix_alloc(2, 2);
  S_inverse_ = gsl_matrix_alloc(2, 2);
  Q_ = gsl_matrix_alloc(2, 2);
  K_ = gsl_matrix_alloc(3, 2);
  I_ = gsl_matrix_alloc(3, 3);
  aux_3x1_ = gsl_matrix_alloc(3, 1);
  aux_3x2_ = gsl_matrix_alloc(3, 2);
  aux_3x3_ = gsl_matrix_alloc(3, 3);
  permutation_ = gsl_permutation_alloc(2);

  gsl_matrix_set_zero(z_predicted_);
  gsl_matrix_set_zero(z_observed_);
  gsl_matrix_set_zero(H_);
  gsl_matrix_set_zero(S_);
  gsl_matrix_set_zero(S_inverse_);
  gsl_matrix_set_zero(Q_);
  gsl_matrix_set_zero(K_);
  gsl_matrix_set_zero(I_);
  gsl_matrix_set_zero(aux_3x1_);
  gsl_matrix_set_zero(aux_3x2_);
  gsl_matrix_set_zero(aux_3x3_);

  gsl_matrix_set(H_, 1, 2, -1.0);

  gsl_matrix_set(Q_, 0, 0, kVarR);
  gsl_matrix_set(Q_, 1, 1, kVarTheta);
}

VisionModel::~VisionModel() {
  gsl_matrix_free(z_predicted_);
  gsl_matrix_free(z_observed_);
  gsl_matrix_free(H_);
  gsl_matrix_free(S_);
  gsl_matrix_free(S_inverse_);
  gsl_matrix_free(Q_);
  gsl_matrix_free(K_);
  gsl_matrix_free(I_);
  gsl_matrix_free(aux_3x1_);
  gsl_matrix_free(aux_3x2_);
  gsl_matrix_free(aux_3x3_);
  gsl_permutation_free(permutation_);
}

void VisionModel::UpdateBelief(player_pose2d *mean, gsl_matrix *covariance) {
  std::vector<player_point_2d> detected_lamps = lamp_centroids_->lamp_centroids();
  if (detected_lamps.empty())
    return;
  std::vector<LampAssociation> associations = AssociateLamps(mean, covariance,
          detected_lamps, mapped_lamps_);

  if (kDebugMode)
    PrintDebugInfo(detected_lamps, associations);

  for (std::vector<LampAssociation>::iterator it = associations.begin();
          it != associations.end(); ++it) {
    // Initializes the observed measurement matrix
    gsl_matrix_set(z_observed_, 0, 0, it->observed_r);
    gsl_matrix_set(z_observed_, 1, 0, it->observed_theta);

    // Initializes the predicted measurement matrix
    gsl_matrix_set(z_predicted_, 0, 0, it->predicted_r);
    gsl_matrix_set(z_predicted_, 1, 0, it->predicted_theta);

    // Initializes matrix H
    gsl_matrix_set(H_, 0, 0, (mean->px - it->mapped_x) / it->predicted_r);
    gsl_matrix_set(H_, 0, 1, (mean->py - it->mapped_y) / it->predicted_r);
    gsl_matrix_set(H_, 1, 0,
            (it->mapped_y - mean->py) / (it->predicted_r * it->predicted_r));
    gsl_matrix_set(H_, 1, 1,
            (mean->px - it->mapped_x) / (it->predicted_r * it->predicted_r));

    // Calculates matrix S
    gsl_blas_dgemm(CblasNoTrans, CblasTrans, 1.0, covariance, H_, 0.0,
            aux_3x2_);
    gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, H_, aux_3x2_, 0.0,
            S_);
    gsl_matrix_add(S_, Q_);

    // Calculates matrix K
    int signum = 0;
    gsl_linalg_LU_decomp(S_, permutation_, &signum);
    gsl_linalg_LU_invert(S_, permutation_, S_inverse_);

    gsl_blas_dgemm(CblasNoTrans, CblasTrans, 1.0, covariance, H_,
            0.0, aux_3x2_);
    gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, aux_3x2_, S_inverse_,
            0.0, K_);

    // Updates the mean
    gsl_matrix_sub(z_observed_, z_predicted_);
    gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, K_, z_observed_, 0.0,
            aux_3x1_);

    mean->px += gsl_matrix_get(aux_3x1_, 0, 0);
    mean->py += gsl_matrix_get(aux_3x1_, 1, 0);
    mean->pa += Localization::AdjustAngle(gsl_matrix_get(aux_3x1_, 2, 0));

    // Updates the covariance
    gsl_matrix_set_identity(I_);
    gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, -1.0, K_, H_, 1.0, I_);
    gsl_blas_dgemm(CblasNoTrans, CblasNoTrans, 1.0, I_, covariance, 0.0,
            aux_3x3_);
    gsl_matrix_memcpy(covariance, aux_3x3_);
  }
}

namespace {
std::vector<LampAssociation> AssociateLamps(const player_pose2d *mean,
        const gsl_matrix *covariance,
        const std::vector<player_point_2d> detected_lamps,
        const std::vector<player_point_2d> mapped_lamps) {
  int lamps = detected_lamps.size();

  // Calculates a search radius for the lamps based on the uncertainty components
  double x_tolerance_squared = 9 * gsl_matrix_get(covariance, 0, 0);
  double y_tolerance_squared = 9 * gsl_matrix_get(covariance, 1, 1);
  double search_radius = sqrt(x_tolerance_squared + y_tolerance_squared)
          + 3 * sqrt(VisionModel::kVarR);
  if (search_radius > 2.0)
    search_radius = 2.0;

  // Converts the detected lamps' coordinates from the robot's frame of reference
  // to the global map's frame of reference
  std::vector<player_point_2d> expected_lamps;
  for (int i = 0; i < lamps; ++i) {
    player_point_2d aux;
    aux.px = mean->px + detected_lamps[i].px * sin(mean->pa)
            + detected_lamps[i].py * cos(mean->pa);
    aux.py = mean->py - detected_lamps[i].px * cos(mean->pa)
            + detected_lamps[i].py * sin(mean->pa);
    expected_lamps.push_back(aux);
  }

  // Makes associations between detected lamps and mapped lamps
  std::vector<LampAssociation> associations;
  std::set<int> associated_mapped_lamps;
  for (int i = 0; i < lamps; ++i) {
    player_point_2d observed = detected_lamps[i];
    player_point_2d expected = expected_lamps[i];
    double min_distance = 100000;
    int min_index = -1;

    // Searches for the mapped lamp closest to the detected lamp
    for (unsigned int j = 0; j < mapped_lamps.size(); ++j) {
      player_point_2d mapped = mapped_lamps[j];
      double dx = mapped.px - expected.px;
      double dy = mapped.py - expected.py;
      if (abs(dx) > search_radius || abs(dy) > search_radius)
        continue;
      double distance = sqrt(dx * dx + dy * dy);
      if (distance < min_distance) {
        min_distance = distance;
        min_index = j;
      }
    }
    if (min_distance > search_radius || min_index == -1)
      continue;

    // Safety device that avoids multiple associations with the same mapped lamp
    if (associated_mapped_lamps.find(min_index)
            != associated_mapped_lamps.end()) {
      printf(
              "[WARNING] Multiple detected lamps associated with the same mapped lamp.\n\n");
      associations.clear();
      return associations;
    }

    // Associates the detected lamp with the mapped lamp found and calculates the
    // observed and predicted values of 'r' and 'theta' of the association.
    LampAssociation aux;
    aux.mapped_x = mapped_lamps[min_index].px;
    aux.mapped_y = mapped_lamps[min_index].py;

    aux.observed_r = sqrt(observed.px * observed.px + observed.py * observed.py);
    aux.observed_theta = atan2(-observed.px, observed.py);

    double dx = mapped_lamps[min_index].px - mean->px;
    double dy = mapped_lamps[min_index].py - mean->py;
    aux.predicted_r = sqrt(dx * dx + dy * dy);
    aux.predicted_theta = Localization::AdjustAngle(atan2(dy, dx) - mean->pa);

    // Filter devised in order to avoid the special case in which a lamp is roughly
    // above the robot and a small position error induces a large orientation error
    if (abs(aux.observed_theta - aux.predicted_theta) > Localization::kPi / 4)
      continue;

    // Numerical correction
    if (aux.observed_r < 1e-6)
      aux.observed_r = 1e-6;
    if (aux.predicted_r < 1e-6)
      aux.predicted_r = 1e-6;

    associations.push_back(aux);
  }

  return associations;
}

void PrintDebugInfo(const std::vector<player_point_2d> detected_lamps,
        const std::vector<LampAssociation> associations) {
  printf("------------------- VISION MODEL DEBUG -------------------\n\n");

  printf("* DETECTED LAMPS *");
  for (std::vector<player_point_2d>::const_iterator it = detected_lamps.begin();
          it != detected_lamps.end(); ++it)
    printf("X = %9.5lf, Y = %9.5lf\n", it->px, it->py);
  printf("\n");

  printf("* ASSOCIATIONS *");
  printf("RATIO: %d/%d\n", (int) associations.size(), (int) detected_lamps.size());
  for (std::vector<LampAssociation>::const_iterator it = associations.begin();
          it != associations.end(); ++it)
    printf("OBS_R: %9.5lf, OBS_THETA: %9.5lf --> PRE_R: %9.5lf, PRE_THETA: %9.5lf\n",
            it->observed_r, it->observed_theta, it->predicted_r, it->predicted_theta);
  printf("\n");

  printf("----------------------------------------------------------\n\n");
}
}

